import numpy as np
import math
import matplotlib.pyplot as plt

N = 5 ** 2  # 绘制点总数
fs = 6
n = [2 * np.pi * fs * t / N for t in range(N)]
x = [int(np.round(np.sin(i) * 1024)) + int(np.round(np.cos(i) * 1024)) * 1j for i in n]
# 频域抽取的基5FFT
loop_num = int(round(math.log(N, 5)))
data = np.zeros((loop_num + 1, N), dtype=np.complex128)
data[0, :] = x

# 基5 FFT 核心算法
for i in range(loop_num):  # 遍历每一层分解
    k = i + 1  # 当前层数对应的基数指数
    for p in range(5 ** i):  # 遍历每组的偏移量
        for j in range(N // (5 ** k)):  # 遍历每组内的元素
            # 第一个分支：直接求和
            data[i + 1][j + 5 * p * (N / (5 ** k))] = (
                data[i, j + 5 * p * (N // (5 ** k))]
                + data[i, j + N // (5 ** k) + 5 * p * (N // (5 ** k))]
                + data[i, j + 2 * N // (5 ** k) + 5 * p * (N // (5 ** k))]
                + data[i, j + 3 * N // (5 ** k) + 5 * p * (N // (5 ** k))]
                + data[i, j + 4 * N // (5 ** k) + 5 * p * (N // (5 ** k))]
            )

            # 第二个分支：引入旋转因子 W_5^1
            data[i + 1][j + N // (5 ** k) + 5 * p * (N // (5 ** k))] = (
                data[i, j + 5 * p * (N // (5 ** k))]
                + data[i, j + N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 1 * 2 * np.pi / 5)
                + data[i, j + 2 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 2 * 2 * np.pi / 5)
                + data[i, j + 3 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 3 * 2 * np.pi / 5)
                + data[i, j + 4 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 4 * 2 * np.pi / 5)
            ) * np.exp(-1j * 2 * j * np.pi * (5 ** i) / N)

            # 第三个分支：引入旋转因子 W_5^2
            data[i + 1][j + 2 * N // (5 ** k) + 5 * p * (N // (5 ** k))] = (
                data[i, j + 5 * p * (N // (5 ** k))]
                + data[i, j + N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 2 * 2 * np.pi / 5)
                + data[i, j + 2 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 4 * 2 * np.pi / 5)
                + data[i, j + 3 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 1 * 2 * np.pi / 5)
                + data[i, j + 4 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 3 * 2 * np.pi / 5)
            ) * np.exp(-1j * 4 * j * np.pi * (5 ** i) / N)

            # 第四个分支：引入旋转因子 W_5^3
            data[i + 1][j + 3 * N // (5 ** k) + 5 * p * (N // (5 ** k))] = (
                data[i, j + 5 * p * (N // (5 ** k))]
                + data[i, j + N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 3 * 2 * np.pi / 5)
                + data[i, j + 2 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 1 * 2 * np.pi / 5)
                + data[i, j + 3 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 4 * 2 * np.pi / 5)
                + data[i, j + 4 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 2 * 2 * np.pi / 5)
            ) * np.exp(-1j * 6 * j * np.pi * (5 ** i) / N)

            # 第五个分支：引入旋转因子 W_5^4
            data[i + 1][j + 4 * N // (5 ** k) + 5 * p * (N // (5 ** k))] = (
                data[i, j + 5 * p * (N // (5 ** k))]
                + data[i, j + N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 4 * 2 * np.pi / 5)
                + data[i, j + 2 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 3 * 2 * np.pi / 5)
                + data[i, j + 3 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 2 * 2 * np.pi / 5)
                + data[i, j + 4 * N // (5 ** k) + 5 * p * (N // (5 ** k))] * np.exp(-1j * 1 * 2 * np.pi / 5)
            ) * np.exp(-1j * 8 * j * np.pi * (5 ** i) / N)

def rev5(k, N):
    if k == 0:
        return 0
    else:
        return (rev5(k // 5, N) // 5) + (k % 5) * (N // 5)


# 输出倒序
fft_out = np.ones_like(data[0, :])
for k in range(N):
    fft_out[rev5(k, N)] = int(data[loop_num,k].real+0.5)+int(data[loop_num,k].imag+0.5)*1j
print(fft_out)

xf = np.fft.fft(x,N)
plt.plot(abs(xf))

fft5_cmp = xf - fft_out

plt.plot(fft5_cmp.imag)
plt.plot(fft5_cmp.real)
plt.show()